import os

import cv2

from show_color import change_color_dark
from show_color import change_color_light

img_seed = 0
txt_seed = 0


def __generate_imgname():
    global img_seed
    img_seed += 1
    return img_seed


def __generate_txtname():
    global txt_seed
    txt_seed += 1
    return txt_seed


def write_txt(path, newname):
    with open(path, 'r', encoding='utf-8') as r:
        with open(newname, 'w', encoding='utf-8') as f:
            txt = r.read(100)
            f.write(txt)
            f.close()
        r.close()
    print('完成写入', newname)


# 改变口罩颜色，分为黑色和白色
def write_imgname(path):
    file_name_black = 'img_add_' + 'black' + str(__generate_imgname()).zfill(4) + '.jpg'
    file_name_white = 'img_add_' + 'white' + str(__generate_imgname()).zfill(4) + '.jpg'
    new_path1 = r'E:\2021homework\cv_masks\yolo\new_masks\black\images'
    new_path2 = r'E:\2021homework\cv_masks\yolo\new_masks\white\images'
    newname1 = new_path1 + os.sep + file_name_black
    newname2 = new_path2 + os.sep + file_name_white
    pic1 = change_color_dark(path)
    if newname1 is not None:
        cv2.imwrite(newname1, pic1)
        print('完成写入', newname1)
    pic2 = change_color_light(path)
    if newname2 is not None:
        cv2.imwrite(newname2, pic2)
        print('完成写入', newname2)


# 改变口罩图片名字，分为黑色和白色
def write_labelname(path):
    file_name_black = 'img_add_' + 'black' + str(__generate_txtname()).zfill(4) + '.txt'
    file_name_white = 'img_add_' + 'white' + str(__generate_txtname()).zfill(4) + '.txt'
    new_path1 = r'E:\2021homework\cv_masks\yolo\new_masks\black\labels'
    new_path2 = r'E:\2021homework\cv_masks\yolo\new_masks\white\labels'
    newname1 = new_path1 + os.sep + file_name_black
    newname2 = new_path2 + os.sep + file_name_white
    if newname1 is not None:
        write_txt(path, newname1)

    if newname2 is not None:
        write_txt(path, newname2)


# 找到各文件夹下图片或者标签
def find_file(path):
    if os.path.exists(path):
        # 判断是文件时，更改名称
        if os.path.isfile(path):
            if path.split('.')[-1] == 'jpg':
                if img_seed < 200:
                    write_imgname(path)
            if path.split('.')[-1] == 'txt':
                if txt_seed < 200:
                    write_labelname(path)

        # 如果路径是文件，进入循环
        else:
            dirs = os.listdir(path)
            for dir in dirs:
                temp_path = path + os.sep + dir  # 这里不能直接用path = path +...，因为for每层递归结束时path保留，相当于path直接加上上一次的后缀
                find_file(temp_path)
                # print(path)
    else:
        return False


if __name__ == '__main__':
    path_img = r'E:\2021homework\cv_masks\yolo\data\mydata\images\val'
    path_txt = r'E:\2021homework\cv_masks\yolo\data\mydata\labels\val'
    # path = r'E:\2021课程\机器视觉与图像测量\大作业\数据集\0003'
    find_file(path_img)
    find_file(path_txt)
